package org.aeon.fasttasks.core.manager;

import jakarta.annotation.Resource;
import org.aeon.fasttasks.core.Task;
import org.aeon.fasttasks.core.annotation.*;
import org.aeon.fasttasks.core.properties.TaskProperties;
import org.aeon.fasttasks.core.recorder.TaskRecorder;
import org.aeon.fasttasks.core.scanner.TaskScanResult;
import org.aeon.fasttasks.core.util.DateUtil;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.NoSuchBeanDefinitionException;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.ApplicationArguments;
import org.springframework.boot.ApplicationRunner;
import org.springframework.context.ApplicationContext;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler;

import java.time.Instant;
import java.util.*;
import java.util.stream.Collectors;
import java.util.stream.Stream;

/**
 * @Date 2024/9/10 0010 15:43
 * @Author aeon
 * @Description 任务调度和执行管理
 */
public class TaskManager implements ApplicationRunner {

    private final static Logger log = LoggerFactory.getLogger(TaskManager.class);

    @Autowired
    private ApplicationContext applicationContext;

    @Resource(name = "batchWorkTasksExecutor")
    private ThreadPoolTaskExecutor batchWorkTasksExecutor;

    @Resource(name = "bossTasksExecutor")
    private ThreadPoolTaskExecutor bossTasksExecutor;

    @Resource(name = "tasksScheduler")
    private ThreadPoolTaskScheduler tasksScheduler;

    @Autowired
    private TaskRecorder taskRecorder;

    @Autowired
    private TaskProperties taskProperties;

    @Override
    public void run(ApplicationArguments args) {
        checkThreadPool();
        bossTasksExecutor.execute(this::arrangeTasks);
    }


    private void arrangeTasks() {
        try {
            List<Task> tasks = getAvailableTasks();
            log.info("共{}个任务", tasks.size());
            scheduleAndExecuteTasks(buildTasks(tasks));
            log.info("所有单次运行任务安排完毕...");
            // 等待所有任务完成
            waitAllTasks(true);
        } catch (Exception e) {
            waitAllTasks(false);
            throw e;
        }
    }

    /**
     * 扫描任务
     * 先查找该类是否注册为bean，
     * 否则通过反射无参构造创建该类
     *
     * @return 可用任务
     */
    private List<Task> getAvailableTasks() {
        return Stream.concat(
                        applicationContext.getBean(TaskScanResult.NAME, TaskScanResult.class).taskClasses().stream(),
                        taskProperties.taskClasses().stream().map(clazzStr -> {
                            try {
                                return Class.forName(clazzStr);
                            } catch (Exception e) {
                                throw new RuntimeException(e);
                            }
                        })
                )
                .distinct()
                .filter(clazz -> {
                    if (!(Task.class.isAssignableFrom(clazz))) {
                        log.warn("{} 任务应当实现 Task 接口", clazz.getName());
                        return false;
                    }
                    return true;
                })
                .map(clazz -> {
                    try {
                        return (Task) applicationContext.getBean(clazz);
                    } catch (NoSuchBeanDefinitionException e) {
                        try {
                            return (Task) clazz.getConstructor().newInstance();
                        } catch (Exception ex) {
                            log.info("任务创建失败：【{}】", clazz.getName());
                            throw new RuntimeException(ex);
                        }
                    } catch (Exception e) {
                        throw new RuntimeException(e);
                    }
                })
                .toList();
    }

    /**
     * 扫描任务类
     *
     * @return 任务类
     * @see TaskScan
     * @see EnableTask
     */
    private Collection<? extends Class<?>> scanTaskClass() {
        return null;
    }

    /**
     * 构建任务
     *
     * @param tasks
     * @return
     */
    private Deque<TaskExpansion> buildTasks(List<Task> tasks) {

        Map<String, Boolean> completedTasks = Objects.nonNull(taskRecorder) ? taskRecorder.readTask() : Map.of();

        Map<Class<? extends Task>, TaskExpansion> taskExpansionMap = tasks
                .stream()
                .filter(task -> {
                    boolean skip = task.getClass().isAnnotationPresent(SkipTask.class) ||
                            (task.isOnlyOnce() && completedTasks.getOrDefault(task.getClass().getName(), false));
                    if (skip) {
                        log.info("任务跳过：【{}】", task.getName());
                    }
                    return !skip;
                })
                .collect(Collectors.toMap(
                        Task::getClass,
                        task -> TaskExpansionFactory.create(task, applicationContext, taskRecorder))

                );

        return taskExpansionMap.values()
                .stream()
                .peek(task -> processDependentTasks(task, taskExpansionMap))
                .sorted(this::compareByPriority)
                .collect(Collectors.toCollection(ArrayDeque::new));

    }

    /**
     * 处理任务依赖关系
     *
     * @param taskExpansion
     * @param taskExpansionMap
     */
    private void processDependentTasks(TaskExpansion taskExpansion,
                                       Map<Class<? extends Task>, TaskExpansion> taskExpansionMap) {
        if (taskExpansion.getTaskContext().getAnnotations().containsKey(DependentTasks.class)) {
            DependentTasks annotation = taskExpansion.getAnnotation(DependentTasks.class);
            Class<? extends Task>[] dependentTasks = annotation.value();
            TaskExpansion[] frontEndTasks = Arrays.stream(dependentTasks)
                    .map(clazz -> {
                        if (taskExpansionMap.containsKey(clazz)) {
                            return taskExpansionMap.get(clazz);
                        }
                        throw new IllegalStateException("类型为 " + clazz.getName() + " 的任务未注册...");
                    })
                    .toArray(TaskExpansion[]::new);
            taskExpansion.getTaskContext().setFrontEndTasks(frontEndTasks);
        }
    }

    /**
     * 按优先级排序任务
     *
     * @param t1
     * @param t2
     * @return
     */
    private int compareByPriority(TaskExpansion t1, TaskExpansion t2) {
        return t2.getTaskContext().getTask().getPriority() - t2.getTaskContext().getTask().getPriority();
    }

    /**
     * 调度和执行任务
     *
     * @param taskExpansions
     * @return
     */
    private void scheduleAndExecuteTasks(Deque<TaskExpansion> taskExpansions) {
        List<TaskExpansion> alreadySchedule = new ArrayList<>();
        while (!taskExpansions.isEmpty()) {
            TaskExpansion task = taskExpansions.peek();

            if (task.isTaskCanStart()) {
                // future确保数据准备方法执行完毕
                task.setFuture(batchWorkTasksExecutor.submit(() -> task.prepareData(false)));
                scheduleTask(taskExpansions, alreadySchedule);
            } else {
                if (Stream.concat(taskExpansions.stream(), alreadySchedule.stream())
                        .anyMatch(t -> task.getTaskContext().getFrontEndTasks().contains(t))) {
                    taskExpansions.addLast(taskExpansions.pop());
                } else {
                    logMissingDependencies(task);
                    taskExpansions.pop();
                }
            }
        }
    }

    /**
     * 安排任务
     *
     * @param pendingTasks
     * @return
     */
    private void scheduleTask(Deque<TaskExpansion> pendingTasks, List<TaskExpansion> alreadySchedule) {
        TaskExpansion task = pendingTasks.pop();
        alreadySchedule.add(task);
        if (task.getTaskContext().isCompleted()) {
            task.getTaskContext().setStart(true);
        } else {
            LoopTask loopTask;
            if (Objects.isNull(task.getAnnotation(LoopTaskPro.class))) {
                loopTask = task.getAnnotation(LoopTask.class);
            } else {
                loopTask = task.getAnnotation(LoopTaskPro.class).async();
            }
            if (Objects.nonNull(loopTask) && loopTask.async() && loopTask.count() > 1) {
                List<TaskExpansion> cloneTasks = new ArrayList<>();
                for (int i = 1; i < loopTask.count(); i++) {
                    TaskExpansion clone = TaskExpansionFactory.clone(task, taskRecorder);
                    cloneTasks.add(clone);
                }
                cloneTasks.add(task);
                cloneTasks.forEach(t -> tasksScheduler.schedule(() -> batchWorkTasksExecutor.execute(t),
                        Instant.now().plus(t.getTaskContext().getSerNum() * loopTask.interval(),
                                DateUtil.toChronoUnit(loopTask.timeUnit()))));
            } else {
                batchWorkTasksExecutor.execute(task);
            }
        }
    }

    /**
     * 记录缺失的依赖
     *
     * @param task
     */
    private void logMissingDependencies(TaskExpansion task) {
        String missingDeps = task.getTaskContext().getFrontEndTasks()
                .stream()
                .map(TaskExpansion::getTaskName)
                .collect(Collectors.joining("】，【"));

        log.error("任务 {} 的前置任务未加载，无法启动已丢弃，请检查：【{}】任务", task.getTaskName(), missingDeps);
    }

    /**
     * 等待所有任务完成
     */
    private void waitAllTasks(boolean isWaitTask) {
        bossTasksExecutor.initiateShutdown();
        while (isWaitTask) {
            try {
                Thread.sleep(500);
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
                break;
            }
            if (tasksScheduler.getScheduledThreadPoolExecutor().getActiveCount() > 0 ||
                    !tasksScheduler.getScheduledThreadPoolExecutor().getQueue().isEmpty()) {
                continue;
            }
            if (batchWorkTasksExecutor.getThreadPoolExecutor().getActiveCount() > 0) {
                continue;
            }
            isWaitTask = false;
        }
        log.info("所有任务执行完毕，程序结束...");
        tasksScheduler.initiateShutdown();
        batchWorkTasksExecutor.initiateShutdown();
    }

    /**
     * 检查线程池大小
     */
    private void checkThreadPool() {
        if (batchWorkTasksExecutor.getCorePoolSize() <= 1) {
            throw new RuntimeException("任务线程池核心线程数应当大于1");
        }
    }

}
